package com.heiye.blog.ai.model.gptsovits;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.heiye.blog.ai.model.openrouter.OpenRouterApi;
import com.heiye.blog.ai.model.vo.SoundGenerationReqVO;
import com.heiye.blog.framework.common.response.Response;
import com.heiye.blog.framework.common.util.JsonUtils;
import lombok.Builder;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.deepseek.api.DeepSeekApi;
import org.springframework.ai.model.ApiKey;
import org.springframework.ai.model.SimpleApiKey;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.util.Assert;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import java.util.Objects;
import java.util.function.Consumer;
import java.util.function.Predicate;

/**
 * @author: heiye
 * @date: 2025/11/20 下午3:32
 * @version: v1.0.0
 * @description: GPT-SoVITS-Api
 */
@Slf4j
@Data
public class GPTSoVITSApi {
    /**
     * 情感模式调用地址
     */
    private final String completionsPath;
    /**
     * 构建 webClient
     */
    private final WebClient webClient;

    /**
     * 生成音频
     *
     * @param options
     * @return
     */
    public Flux<DataBuffer> generation(GPTSoVITSOptions options) {
        // 创建请求体
        Request request = createRequest(options);

        log.info("GPT-SoVITS 请求参数：{}", JsonUtils.toJsonString(request));


        return this.webClient.post()
                .uri(this.getEndpoint())
                .bodyValue(request)
                .retrieve()
                // 流式处理
                .bodyToFlux(DataBuffer.class);
    }

    private String getEndpoint() {
        return this.completionsPath;
    }

    public static Builder builder() {
        return new Builder();
    }
    /**
     * 创建请求体
     *
     * @param options
     * @return
     */
    public Request createRequest(GPTSoVITSOptions options) {

        // 先构建 otherParams
        OtherParams otherParams = OtherParams.builder()
                .appKey(options.getAppKey())
                .textLang(options.getTextLang())
                .promptLang(options.getPromptLang())
                .emotion(options.getEmotion())
                .topK(options.getTopK())
                .topP(options.getTopP())
                .temperature(options.getTemperature())
                .textSplitMethod(options.getTextSplitMethod())
                .batchSize(options.getBatchSize())
                .batchThreshold(options.getBatchThreshold())
                .splitBucket(options.getSplitBucket())
                .fragmentInterval(options.getFragmentInterval())
                .parallelInfer(options.getParallelInfer())
                .repetitionPenalty(options.getRepetitionPenalty())
                .sampleSteps(options.getSampleSteps())
                .ifSr(options.getIfSr())
                .seed(options.getSeed())
                .build();

        // 最后再构建请求体
        return Request.builder()
                .model(options.getModel())
                .input(options.getText())
                .voice(options.getRoleName())
                .speed(options.getSpeed())
                .responseFormat(options.getMediaType())
                .otherParams(otherParams)
                .build();
    }

    /**
     * 构造函数
     *
     * @param baseUrl
     * @param completionsPath
     * @param headers
     * @param webClientBuilder
     */
    public GPTSoVITSApi(String baseUrl, String completionsPath, MultiValueMap<String, String> headers, WebClient.Builder webClientBuilder) {
        this.completionsPath = completionsPath;

        Consumer<HttpHeaders> finalHeaders = (h) -> {
            h.setContentType(MediaType.APPLICATION_JSON);
            h.addAll(headers);
        };

        this.webClient = webClientBuilder.baseUrl(baseUrl).defaultHeaders(finalHeaders).build();
    }

    /**
     * 创建 Builder 模式
     */
    public static class Builder {
        /**
         * 默认调用地址
         */
        private String baseUrl = "http://127.0.0.1:8000/";
        private MultiValueMap<String, String> headers = new LinkedMultiValueMap<>();
        private String completionsPath = "/v1/audio/speech";
        private WebClient.Builder webClientBuilder = WebClient.builder();

        public Builder baseUrl(String baseUrl) {
            this.baseUrl = baseUrl;
            return this;
        }


        public Builder headers(MultiValueMap<String, String> headers) {
            this.headers = headers;
            return this;
        }

        public Builder webClientBuilder(WebClient.Builder webClientBuilder) {
            this.webClientBuilder = webClientBuilder;
            return this;
        }


        public Builder completionsPath(String completionsPath) {
            this.completionsPath = completionsPath;
            return this;
        }

        public GPTSoVITSApi build() {
            return new GPTSoVITSApi(this.baseUrl, this.completionsPath, this.headers, this.webClientBuilder);
        }
    }

    // GPT-SoVITS 请求参数
    @lombok.Builder
    @JsonInclude(JsonInclude.Include.NON_NULL)
    public record Request(
            // 模型名称
            @JsonProperty("model")
            String model,

            // 待合成的文本内容
            @JsonProperty("input")
            String input,

            // 角色模型名称
            @JsonProperty("voice")
            String voice,

            // 语速
            @JsonProperty("speed")
            Integer speed,

            // 媒体文件类型
            @JsonProperty("response_format")
            String responseFormat,

            // 其他参数
            @JsonProperty("other_params")
            OtherParams otherParams

    ) {

        public Request(@JsonProperty("model")
                       String model, @JsonProperty("input")
                       String input, @JsonProperty("voice")
                       String voice, @JsonProperty("speed")
                       Integer speed, @JsonProperty("response_format")
                       String responseFormat, @JsonProperty("other_params")
                       OtherParams otherParams) {
            this.model = model;
            this.input = input;
            this.voice = voice;
            this.speed = speed;
            this.responseFormat = responseFormat;
            this.otherParams = otherParams;
        }
    }

    // GPT-SoVITS 其他参数
    @lombok.Builder
    @JsonInclude(JsonInclude.Include.NON_NULL)
    public record OtherParams(
            // 应用密钥
            @JsonProperty("app_key")
            String appKey,

            // 文本语言
            @JsonProperty("text_lang")
            String textLang,

            // 提示文本语言
            @JsonProperty("prompt_lang")
            String promptLang,

            // 情感参数
            @JsonProperty("emotion")
            String emotion,

            // 采样时选择的top k个最可能的token
            @JsonProperty("top_k")
            Integer topK,

            // 采样时累积概率阈值
            @JsonProperty("top_p")
            Integer topP,

            // 采样温度参数，控制生成随机性
            @JsonProperty("temperature")
            Integer temperature,

            // 文本分割方法
            @JsonProperty("text_split_method")
            String textSplitMethod,

            // 批处理大小
            @JsonProperty("batch_size")
            Integer batchSize,

            // 批处理阈值
            @JsonProperty("batch_threshold")
            Double batchThreshold,

            // 是否分割桶
            @JsonProperty("split_bucket")
            Boolean splitBucket,

            // 片段间隔时间
            @JsonProperty("fragment_interval")
            Double fragmentInterval,

            // 是否并行推理
            @JsonProperty("parallel_infer")
            Boolean parallelInfer,

            // 重复惩罚系数
            @JsonProperty("repetition_penalty")
            Double repetitionPenalty,

            // 采样步数
            @JsonProperty("sample_steps")
            Integer sampleSteps,

            // 是否进行语音识别
            @JsonProperty("if_sr")
            Boolean ifSr,

            // 随机种子
            @JsonProperty("seed")
            Integer seed
    ) {
    }
}
